package com.pacvue.handler;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class LocalRateLimiterAcquireTokenHandler extends AbstractRateLimiterAcquireTokenHandler{

    private static final Map<String, LocalBucket> LOCAL_CACHE = new ConcurrentHashMap<>();
    private static final Map<String, ReentrantLock> LOCK_MAP = new ConcurrentHashMap<>();
    private static final ScheduledExecutorService CLEANUP_EXECUTOR = Executors.newSingleThreadScheduledExecutor();
    
    // 缓存清理间隔（分钟）
    private static final int CLEANUP_INTERVAL_MINUTES = 30;
    // 桶过期时间（分钟）
    private static final int BUCKET_EXPIRE_MINUTES = 60;

    static {
        // 启动定期清理任务
        CLEANUP_EXECUTOR.scheduleAtFixedRate(
            LocalRateLimiterAcquireTokenHandler::cleanupExpiredBuckets,
            CLEANUP_INTERVAL_MINUTES,
            CLEANUP_INTERVAL_MINUTES,
            TimeUnit.MINUTES
        );
        
        // 添加JVM关闭钩子
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            CLEANUP_EXECUTOR.shutdown();
            try {
                if (!CLEANUP_EXECUTOR.awaitTermination(5, TimeUnit.SECONDS)) {
                    CLEANUP_EXECUTOR.shutdownNow();
                }
            } catch (InterruptedException e) {
                CLEANUP_EXECUTOR.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }));
    }

    public LocalRateLimiterAcquireTokenHandler() {}

    @Override
    public boolean tryAcquire(String key, int capacity, double rate) {
        ReentrantLock lock = LOCK_MAP.computeIfAbsent(key, k -> new ReentrantLock());
        
        try {
            lock.lock();
            long now = System.currentTimeMillis();

            // 初始化
            LocalBucket bucket = LOCAL_CACHE.get(key);
            if (bucket == null) {
                bucket = new LocalBucket();
                bucket.setTokens(capacity);
                bucket.setLastRefill(now);
                bucket.setLastAccess(now);
                LOCAL_CACHE.put(key, bucket);
            } else {
                bucket.setLastAccess(now);
            }

            double tokens = bucket.getTokens();
            long lastRefill = bucket.getLastRefill();

            // 补充令牌
            long delta = Math.max(0, now - lastRefill);
            double added_tokens = delta * rate / 1000.0;
            tokens = Math.min(capacity, tokens + added_tokens);
            lastRefill = now;

            // 判断是否可以获取令牌
            boolean flag = false;
            if (tokens >= 1) {
                tokens -= 1;
                flag = true;
            }

            // 更新缓存
            bucket.setTokens(tokens);
            bucket.setLastRefill(lastRefill);
            
            return flag;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 清理过期的桶
     */
    private static void cleanupExpiredBuckets() {
        try {
            long now = System.currentTimeMillis();
            long expireTime = now - (BUCKET_EXPIRE_MINUTES * 60 * 1000L);
            
            LOCAL_CACHE.entrySet().removeIf(entry -> {
                LocalBucket bucket = entry.getValue();
                boolean expired = bucket.getLastAccess() < expireTime;
                if (expired) {
                    LOCK_MAP.remove(entry.getKey());
                    if (log.isDebugEnabled()) {
                        log.debug("清理过期的限流桶: {}", entry.getKey());
                    }
                }
                return expired;
            });
            
            if (log.isDebugEnabled()) {
                log.debug("缓存清理完成，当前桶数量: {}", LOCAL_CACHE.size());
            }
        } catch (Exception e) {
            log.error("清理过期桶时发生异常", e);
        }
    }

    @Data
    public static class LocalBucket {
        private double tokens;              // 当前令牌数
        private long lastRefill;            // 上次补充时间（毫秒）
        private long lastAccess;            // 上次访问时间（毫秒）
    }
}
